import torch
import torch.nn as nn
import torch.optim as optim
import sys
import os

# 添加模块路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from moudle import (
    create_data_loaders, 
    Trainer, 
    setup_device, 
    load_label_mapping,
    TensorBoardLogger,
    YAMLConfig,
    load_model
)


class VGG19(nn.Module):
    """VGG19模型实现"""
    
    def __init__(self, num_classes=100, dropout=0.5):
        super(VGG19, self).__init__()
        
        # VGG19的特征提取层配置
        # 'M' 表示 MaxPooling
        cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 
               512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
        
        self.features = self._make_layers(cfg)
        
        # 自适应平均池化，将特征图大小统一为 7x7
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        
        # 分类器
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(dropout),
            nn.Linear(4096, num_classes),
        )
        
        # 初始化权重
        self._initialize_weights()
    
    def _make_layers(self, cfg, batch_norm=True):
        """构建VGG的卷积层"""
        layers = []
        in_channels = 3
        
        for v in cfg:
            if v == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
                if batch_norm:
                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
                else:
                    layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = v
        
        return nn.Sequential(*layers)
    
    def _initialize_weights(self):
        """初始化权重"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


def create_vgg_config():
    """创建VGG19的默认配置"""
    config = YAMLConfig()
    
    # 更新VGG19专用配置
    vgg_config = {
        # 数据配置
        'data_dir': '../dataset/Mini-ImageNet-Dataset',
        'label_file': '../dataset/label.json',
        'batch_size': 128,  # VGG通常使用较小的batch size
        'num_workers': 0,
        'pin_memory': True,
        'img_size': 224,
        'num_classes': 100,
        
        # 数据增强配置
        'normalize_mean': [0.485, 0.456, 0.406],
        'normalize_std': [0.229, 0.224, 0.225],
        
        # 训练配置
        'num_epochs': 50,
        'lr': 0.001,  # VGG通常使用较小的学习率
        'momentum': 0.9,
        'weight_decay': 1e-4,
        'print_freq': 50,
        
        # 学习率调度器
        'scheduler_type': 'StepLR',
        'step_size': 15,  # 每15个epoch学习率衰减
        'gamma': 0.1,
        
        # 路径配置
        'log_dir': './logs/vgg19',
        'model_save_path': './model/vgg19_model.pth',
        'save_best': True,
        
        # 模型配置
        'model_name': 'VGG19',
        'dropout': 0.5,
    }
    
    config.update(vgg_config)
    return config


def main():
    print("=== VGG19 训练 ===")
    
    # 1. 创建配置
    config = create_vgg_config()
    
    # 检查是否存在YAML配置文件，如果存在则加载
    config_file = os.path.join(os.path.dirname(__file__), 'config.yaml')
    if os.path.exists(config_file):
        print(f"发现配置文件 {config_file}，加载中...")
        config = YAMLConfig(config_file)
    else:
        print("使用默认配置")
        # 保存默认配置到文件
        config.save_config(config_file)
        print(f"默认配置已保存到 {config_file}")
    
    # 打印配置信息
    print("\n=== 当前配置 ===")
    print(f"模型: {config.get('model_name')}")
    print(f"批次大小: {config.get('batch_size')}")
    print(f"学习率: {config.get('lr')}")
    print(f"训练轮数: {config.get('num_epochs')}")
    print(f"图像尺寸: {config.get('img_size')}")
    print(f"Dropout: {config.get('dropout')}")
    print("=" * 20)
    
    # 2. 设置设备
    device = setup_device()
    
    # 3. 创建数据加载器
    print("\n加载数据集...")
    train_loader, val_loader, test_loader, train_dataset = create_data_loaders(
        data_dir=config.get('data_dir'),
        batch_size=config.get('batch_size'),
        num_workers=config.get('num_workers'),
        pin_memory=config.get('pin_memory'),
        img_size=config.get('img_size'),
        mean=config.get('normalize_mean'),
        std=config.get('normalize_std')
    )
    
    print(f"训练集大小: {len(train_dataset)}")
    print(f"训练批次数: {len(train_loader)}")
    print(f"验证批次数: {len(val_loader)}")
    print(f"测试批次数: {len(test_loader)}")
    
    # 4. 加载标签映射
    label_mapping = load_label_mapping(config.get('label_file'))
    
    # 5. 创建模型
    print("\n创建VGG19模型...")
    model = VGG19(
        num_classes=config.get('num_classes'),
        dropout=config.get('dropout')
    )
    
    # 6. 创建损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(
        model.parameters(), 
        lr=config.get('lr'), 
        momentum=config.get('momentum'), 
        weight_decay=config.get('weight_decay')
    )
    
    # 7. 创建学习率调度器
    if config.get('scheduler_type') == 'StepLR':
        scheduler = optim.lr_scheduler.StepLR(
            optimizer, 
            step_size=config.get('step_size'), 
            gamma=config.get('gamma')
        )
    else:
        scheduler = None
    
    # 8. 创建TensorBoard记录器
    tensorboard_logger = TensorBoardLogger(
        log_dir=config.get('log_dir'),
        model_name=config.get('model_name', 'vgg19'),
        comment='training'
    )
    
    # 9. 创建训练器
    trainer = Trainer(
        model=model,
        device=device,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=scheduler,
        tensorboard_logger=tensorboard_logger,
        model_save_path=config.get('model_save_path'),
        print_freq=config.get('print_freq')
    )
    
    # 10. 开始训练
    training_history = trainer.train(
        train_loader=train_loader,
        val_loader=val_loader,
        num_epochs=config.get('num_epochs'),
        save_best=config.get('save_best')
    )
    
    # 11. 测试模型
    print('\n加载最佳模型进行测试...')
    best_model_path = config.get('model_save_path').replace('.pth', '_best.pth')
    model = load_model(model, best_model_path, device)
    
    # 重新创建测试用的训练器
    test_trainer = Trainer(model=model, device=device)
    test_accuracy = test_trainer.test(
        test_loader=test_loader,
        train_dataset=train_dataset,
        label_mapping=label_mapping,
        num_classes=config.get('num_classes')
    )
    
    # 12. 记录最终结果
    tensorboard_logger.log_test_results(test_accuracy)
    tensorboard_logger.log_final_summary(
        training_time=training_history['training_time'],
        best_val_acc=training_history['best_val_acc'],
        test_acc=test_accuracy,
        model_path=config.get('model_save_path')
    )
    
    # 13. 保存最终配置
    final_config_path = os.path.join(config.get('log_dir'), 'final_config.yaml')
    config.save_config(final_config_path)
    print(f"\n最终配置已保存至: {final_config_path}")
    
    # 14. 打印最终结果
    print(f'\n=== VGG19 训练总结 ===')
    print(f'训练时间: {training_history["training_time"]:.2f} 秒')
    print(f'最佳验证准确率: {training_history["best_val_acc"]:.2f}%')
    print(f'最终测试准确率: {test_accuracy:.2f}%')
    print(f'模型已保存至: {config.get("model_save_path")}')
    print('=' * 30)
    
    # 15. 关闭TensorBoard并打印命令
    tensorboard_logger.print_tensorboard_command()
    tensorboard_logger.close()


if __name__ == '__main__':
    main()
